package hmm;

/**
 * Implements the Viterby algorithm
 * 
 * @author Amit Yungman
 */
public class Viterby {
	private Hmm hmm;	// HMM according to which we tag the sentence
	private double[][] delta; //The partial probability matrix
	private int[] y; //The best path
	private int[] k; //The words' indices in the HMM
  private Double[][][] ndelta; //The partial probability matrix
  private Integer[][][] nys; //The paths
  private Integer[] nk; //The words' indices in the HMM
	
	/**
	 * @param _hmm - The HMM to work on
	 */
	public Viterby(Hmm _hmm) {
		hmm = _hmm;
	}
  
  /**
   * Calculates the top possible tagging for the sentence in the HMM
   * @param sentence - The sentence to process - <b>contains only base words!</b>
   * @return An array of the top result of the viterby algorithm
   */
  public ViterbyResult calculateViterby(String[] sentence) {
    initializeValues(sentence);
    for(int t = 1 ; t < sentence.length ; t++) {
      for(int i = 0 ; i < delta.length ; i++)
        setMax(i,t);
      y[t] = maxDelta(t);
    }
    return new ViterbyResult(processPath(), delta[y[sentence.length-1]][sentence.length-1] );
  }

  /**
	 * Calculates the top n possible tagging for the sentence in the HMM
	 * @param sentence - The sentence to process
	 * @param n - The number of top tagging to calculate
	 * @return An array of the n top results of the viterby algorithm
	 */
	public ViterbyResult[] calculateViterbyN(String[] sentence, int n) {
	  initializeValues(sentence,n);
    Integer[] maxStates = new Integer[n];
    for(int t = 1 ; t < sentence.length ; t++) {
      for(int i = 0 ; i < ndelta[0].length ; i++) {
        maxStates = setMax(i,t,n);
        for(int j = 0 ; j < n ; j++)
          nys[j][i][t] = maxStates[j];
      }
    }

    ViterbyResult[] results = new ViterbyResult[n];
    Integer[] maxI = new Integer[n];
    Integer[] maxN = new Integer[n];
    Double[] max = new Double[n];
    int index;
    for(int i = 0 ; i < n ; i++)
      max[i] = -1.0;
    for(int j = 0 ; j < n ; j++)
      for(int i = 0 ; i < hmm.getAllTags().size() ; i++) {
        index = replaceIfGreater(ndelta[j][i][sentence.length-1], max);
        if(index >= 0) {
          replaceArrayValue(maxI, i, index);
          replaceArrayValue(maxN, j, index);
        }
      }
    for(int i = 0 ; i < n ; i++)
      results[i] = new ViterbyResult(processPathN(maxN[i],maxI[i],sentence.length-1,sentence.length), ndelta[maxN[i]][maxI[i]][sentence.length-1]);
    return results;
	}
  
  /************ Private methods ****************/

	/**
	 * Back tracks the path to a certain location in the delta matrix
	 * @param n - The n index of the matrix
	 * @param i - The tag index of the matrix
	 * @param t - The words in the matrix (time)
	 * @param N - The size of the return array
	 */
  private String[] processPathN(Integer n, Integer i, Integer t, int N) {
    int[] $ = new int[N];
    $[N-1] = i;
    if(N > 1) {
      $[N-2] = nys[n][i][t];
      for(int j = 3 ; j <= N ; j++) {
        $[N-j] = nys[0][$[N-j+1]][t-j+2];
      }
    }
    return processPath($);
  }
  

  /**
   * Initializes the values of the calculation
   * @param sentence - The sentence to process
   */
  private void initializeValues(String[] sentence) {
    final int size = sentence.length;
    k = new int[size];
    for(int i = 0 ; i < size ; i++)
      k[i] = hmm.getAllBaseWords().indexOf(sentence[i]);
    
    delta = new double[hmm.getAllTags().size()][size];
    y = new int[size];
    for(int i = 0 ; i < delta.length ; i++)
      delta[i][0] = hmm.getPi()[i] * hmm.getB()[i][k[0]];
    y[0] = maxDelta(0);
  }

  /**
   * Initializes the values of the calculation
   * @param sentence - The sentence to process
   * @param n - The number of tagging
   */
  private void initializeValues(String[] sentence, int n) {
    final int size = sentence.length;
    nk = new Integer[size];
    for(int i = 0 ; i < size ; i++)
        nk[i] = hmm.getAllBaseWords().indexOf(sentence[i]);
    
    ndelta = new Double[n][hmm.getAllTags().size()][size];
    nys = new Integer[n][hmm.getAllTags().size()][size];
    for(int i = 0 ; i < ndelta[0].length ; i++) {
      for(int j = 0 ; j < n ; j++) {
        ndelta[j][i][0] = hmm.getPi()[i] * hmm.getB()[i][nk[0]];
        nys[j][i][0] = i;
        for(int h = 1 ; h < size ; h++) {
          ndelta[j][i][h] = -1.0;
          nys[j][i][h] = -1;
        }
      }
    }
  }

  /**
   * Calculates the maximal path probability to this state at this time
   * @param t - The "time"
   * @param i - The state
   */
  private void setMax(int i, int t) {
    double curr, max = -1;
    for(int j = 0 ; j < delta.length ; j++)
      if((curr = delta[j][t-1] * hmm.getA()[j][i] * hmm.getB()[i][k[t]]) >= max)
        max = curr;
    delta[i][t] = max;
  }

  /**
   * Calculates the n maximal path probabilities to this state at this time
   * @param t - The "time"
   * @param i - The state
   * @param n - The number of top values to process
   */
  private Integer[] setMax(int i, int t, int n) {
    Double[] max = new Double[n];
    Integer[] states = new Integer[n];
    int index;
    for(int h = 0 ; h < n ; h++)
      max[h] = -1.0;
    for(int j = 0 ; j < ndelta[0].length ; j++) {
      index = replaceIfGreater(ndelta[0][j][t-1] * hmm.getA()[j][i] * hmm.getB()[i][nk[t]],max);
      if(index >= 0)
        replaceArrayValue(states, j, index);
    }
    for(int h = 0 ; h < n ; h++)
      ndelta[h][i][t] = max[h];
    return states;
  }
  
  /**
   * Returns the state with the maximal value at t time
   * @param t - The time
   * @return The state with the maximal value at t time
   */
  private int maxDelta(int t) {
    double max = -1;
    int state = -1;
    for(int i = 0 ; i < delta.length ; i++)
      if(delta[i][t] >= max) {
        max = delta[i][t];
        state = i;
      }
    return state;
  }


  /**
   * Process a path into a string path
   * @return The path in y with its tag values
   */
  private String[] processPath() {
    String[] stringPath = new String[y.length];
    for(int i = 0 ; i < y.length ; i++)
      stringPath[i] = hmm.getAllTags().get(y[i]);
    return stringPath;
  }
  
  /**
   * Process a path into a string path
   * @param path - The path to process
   * @return The path in ny[j] with its tag values
   */
  private String[] processPath(int[] path) {
    String[] stringPath = new String[path.length];
    for(int i = 0 ; i < path.length ; i++)
      stringPath[i] = hmm.getAllTags().get(path[i]);
    return stringPath;
  }

  /**
   * Checks if the number if greater from some of the values in the array.
   * If so, adds the number to the array, sorts it, and drops the minimal 
   * number from the array.
   * @param num - The number
   * @param sortedNums - The array to go through
   * @return -1 if didn't replace anything. The index of the new value if 
   * inserted it.
   */
  private int replaceIfGreater(double num, Double[] sortedNums) {
    for(int j = 0 ; j < sortedNums.length ; j++)
      if(num > sortedNums[j]) {
        replaceArrayValue(sortedNums, num, j);
        return j;
      }
    return -1;
  }
  
  /**
   * Replaces the i-th element in the array withnewVal, and pushes all the
   * elements that follow up one step
   * @param array - The array of numbers
   * @param newVal - The new value to insert
   * @param i - The location to insert to
   */
  private void replaceArrayValue(Integer[] array, int newVal, int i) {
    Integer[] copy = array.clone();
    int locationChange = 0;
    for(int j = 0 ; j < array.length ; j++) {
      if(j == i) {
        array[j] = newVal;
        locationChange = -1;
      } 
      else
        array[j] = copy[j+locationChange];
    }
  }

  /**
   * Replaces the i-th element in the array withnewVal, and pushes all the
   * elements that follow up one step
   * @param array - The array of numbers
   * @param newVal - The new value to insert
   * @param i - The location to insert to
   */
  private void replaceArrayValue(Double[] array, double newVal, int i) {
    Double[] copy = array.clone();
    int locationChange = 0;
    for(int j = 0 ; j < array.length ; j++) {
      if(j == i) {
        array[j] = newVal;
        locationChange = -1;
      } 
      else
        array[j] = copy[j+locationChange];
    }
  }
}
